package ffte.core.streamed

import spinal.core._
import spinal.lib._

import ffte.types.VecComplex
import ffte.property.{getFFTIW,getFFTOW}
import ffte.misc
import ffte.core.{samewidthed}

/**
 * Streamed interleaver factory object
 *
 * Provides factory methods for creating optimized data interleavers for different
 * transform size configurations. The interleaver handles the complex data reordering
 * required between FFT stages in streamed implementations.
 *
 * Interleaving Strategy:
 * - Square matrices (first == second): Use element interleaver for efficiency
 * - Special cases (first < second, second % first == 0): Use optimized permutation
 * - General case: Use general-purpose permutation interleaver
 *
 * Data Reordering Purpose:
 * - PSB decomposition: Chinese Remainder Theorem-based index mapping
 * - MSB decomposition: Cooley-Tukey-based index permutation
 * - Required between FFT stages for correct data flow
 *
 * @param dW Data width for complex numbers
 * @param first First dimension size
 * @param second Second dimension size
 * @return Optimized streamed interleaver component
 */
object SInterleaver {
    /**
     * Create optimized data interleaver
     *
     * Selects the most appropriate interleaver implementation based on
     * the transform size characteristics. This optimization provides
     * significant resource savings for common transform configurations.
     *
     * @param dW Data width for complex numbers
     * @param first First dimension size
     * @param second Second dimension size
     * @return Optimized interleaver implementation
     */
    def apply(dW: Int, first: Int, second: Int): streamedComponents =
        if (first == second) (new eInterleaver(dW, first))
        else if (first < second && (second % first) == 0) (new pInterleaver(dW, first, second))
        else (new pInterleaver(dW, first, second))

    /**
     * Hold stream data when valid
     *
     * Utility function that holds stream data in a register when the stream
     * is valid. This is used for timing optimization in the interleaver
     * pipeline to maintain data synchronization.
     *
     * @param x Input stream with generic data type
     * @return Held data when stream is valid
     */
    def hold[T <: Data](x: Stream[T]): T = {
        val ha = new Area {
            val r = Reg(x.payloadType())
            r := x.payload
        }
        (x.valid) ? x.payload | ha.r
    }
}

/**
 * General-purpose permutation interleaver
 *
 * This class implements a general-purpose data interleaver for arbitrary
 * transform dimensions using a ping-pong memory buffer scheme. The interleaver
 * handles complex data reordering required between FFT stages with configurable
 * memory usage and timing characteristics.
 *
 * Architecture Strategy:
 * - Ping-pong memory buffer for efficient data reordering
 * - Dual counter scheme for read/write address generation
 * - Configurable latency based on transform size
 * - Streaming interface with flow control
 *
 * Data Reordering Logic:
 * - Write phase: Sequential input storage in row-major order
 * - Read phase: Column-major output for FFT stage reordering
 * - Memory management: Efficient address calculation for permutations
 *
 * Performance Characteristics:
 * - Supports arbitrary transform dimensions
 * - Configurable latency: first*second + 1 cycles
 * - Memory usage: 2^(fAW+sAW+1) complex words
 * - Efficient for medium to large transforms
 *
 * @param dW Data width for complex numbers
 * @param first First dimension size
 * @param second Second dimension size
 */
sealed class pInterleaver(val dW: Int, first: Int, second: Int) extends streamedComponents with samewidthed {
    /** Standard streaming I/O interface */
    val io = iogen

    /** Address width calculations for memory sizing */
    val fAW = log2Up(first)
    val sAW = log2Up(second)

    /**
     * Ping-pong memory buffer
     *
     * Stores intermediate data during the reordering process.
     * Memory size is calculated to handle the worst-case data
     * reordering scenario with sufficient capacity.
     */
    val mem = Mem(VecComplex(G, dW), 1 << (fAW + sAW + 1))

    /** Flow control: Direct backpressure propagation */
    io.d.ready := io.q.ready
    io.q.valid := io.d.fire
    val fire = io.d.fire

    /**
     * Total interleaver delay
     *
     * Includes memory fill time plus one cycle for output generation.
     * The delay scales with the total number of elements.
     */
    def delay = first * second + 1

    /**
     * Start signal for output phase
     *
     * Generates a pulse when the memory is filled and output
     * phase can begin. This ensures proper synchronization
     * between write and read phases.
     */
    val start = misc.delay(delay - 1, first * second).delayPulse(io.d.payload.last, fire)

    /**
     * Read phase counters
     *
     * rfCnt: First dimension read counter (0 to first-1)
     * rsCnt: Second dimension read counter (0 to second-1)
     * rpp: Ping-pong phase counter (0/1)
     */
    val rfCnt = Counter(first, inc = fire)
    val rsCnt = Counter(second, inc = rfCnt.willOverflow)
    val rpp = Counter(2, inc = rsCnt.willOverflow)

    /**
     * Write phase counters
     *
     * wsCnt: Second dimension write counter (0 to second-1)
     * wfCnt: First dimension write counter (0 to first-1)
     * wpp: Write phase ping-pong control
     */
    val wsCnt = Counter(second, inc = fire)
    val wfCnt = Counter(first, inc = wsCnt.willOverflow)
    val wpp = RegInit(U(0, 1 bits))

    /**
     * Counter reset logic
     *
     * Resets all counters when input burst completes to
     * prepare for the next interleaving operation.
     */
    when(fire) {
        when(io.d.payload.last) {
            rfCnt.clear
            rsCnt.clear
        }
        when(rsCnt.willOverflowIfInc && rfCnt.willOverflowIfInc) {
            wfCnt.clear
            wsCnt.clear
            wpp := rpp.value
        }
    }
    mem.write(
        address = rpp @@ rsCnt @@ rfCnt,
        data    = io.d.payload.fragment,
        enable  = fire
    )
    io.q.payload.fragment := mem.readSync(
        address = wpp @@ wsCnt @@ wfCnt,
        enable  = fire
    )    
    io.q.payload.last := RegNextWhen(start,fire)

}

/**
 * Element interleaver for square matrix transforms
 *
 * This class implements a specialized data interleaver optimized for square
 * matrix transformations (N × N). It provides efficient element-wise
 * data reordering using a transposition-based approach.
 *
 * Architecture Features:
 * - Ping-pong memory buffer for transposition operation
 * - Element-wise transposition for square matrices
 * - Optimized address calculation for N×N matrices
 * - Configurable latency: N² + 1 cycles
 *
 * Transposition Strategy:
 * - Write phase: Row-major storage (standard order)
 * - Read phase: Column-major output (transposed order)
 * - Memory management: Efficient ping-pong buffer usage
 *
 * Optimization:
 * - Specialized address generation for square matrices
 * - Reduced address width compared to general case
 * - Efficient element-wise transposition logic
 *
 * Performance Characteristics:
 * - Optimized for square transform sizes
 * - Fixed delay: N² + 1 cycles for N×N transform
 * - Memory usage: 2^(2*AW) complex words
 * - Most efficient for square matrix FFT stages
 *
 * @param dW Data width for complex numbers
 * @param N Transform size (square matrix dimension)
 */
sealed class eInterleaver(val dW:Int,N:Int) extends streamedComponents with samewidthed {
    /** Standard streaming I/O interface */
    val io = iogen

    /** Flow control: Direct backpressure propagation */
    io.d.ready := io.q.ready
    io.q.valid := io.d.fire
    val fire = io.d.fire

    /** Address width for N×N matrix indexing */
    val AW = log2Up(N)

    /**
     * Ping-pong memory buffer for N×N transposition
     *
     * Stores intermediate data during the element-wise transposition.
     * Memory size optimized for square matrix operations.
     */
    val mem = Mem(VecComplex(G,dW),1<<(2*AW))

    /**
     * Total interleaver delay for N×N transposition
     *
     * Includes memory fill time for N×N elements plus one cycle
     * for output generation.
     */
    def delay = N*N+1

    /**
     * Start signal for output phase
     *
     * Generates a pulse when the N×N memory is filled and output
     * phase can begin with proper synchronization.
     */
    val start = misc.delay(delay-1,N).delayPulse(io.d.payload.last,fire)

    /**
     * Read phase counters for transposed access
     *
     * rfCnt: Row counter for read operations (0 to N-1)
     * rsCnt: Column counter for read operations (0 to N-1)
     * rpp: Ping-pong phase counter (0/1)
     */
    val rfCnt = Counter(N,inc=fire)
    val rsCnt = Counter(N,inc=rfCnt.willOverflow)
    val rpp   = Counter(2,inc=rsCnt.willOverflow)

    /**
     * Write phase counters for row-major storage
     *
     * wsCnt: Column counter for write operations (0 to N-1)
     * wfCnt: Row counter for write operations (0 to N-1)
     * wpp: Write phase ping-pong control
     */
    val wsCnt = Counter(N,inc=fire)
    val wfCnt = Counter(N,inc=wsCnt.willOverflow)
    val wpp   = RegInit(U(0, 1 bits))

    /**
     * Counter reset and synchronization logic
     *
     * Manages counter resets and ping-pong buffer switching
     * when input bursts complete and output phases synchronize.
     */
    when(fire) {
        when(io.d.payload.last) {
            rfCnt.clear
            rsCnt.clear
        }
        when(rsCnt.willOverflowIfInc && rfCnt.willOverflowIfInc) {
            wfCnt.clear
            wsCnt.clear
            wpp := rpp.value
        }
    }

    /**
     * Memory write operation (row-major storage)
     *
     * Writes input data in row-major order for N×N matrix.
     * Address pattern depends on ping-pong phase for transposition.
     */
    mem.write(
        address = ((rpp===0)? (rsCnt @@ rfCnt) | (rfCnt @@ rsCnt)),
        data    = io.d.payload.fragment,
        enable  = fire
    )

    /**
     * Memory read operation (column-major output)
     *
     * Reads data in column-major order for transposed output.
     * Address pattern creates the N×N matrix transposition.
     */
    io.q.payload.fragment := mem.readSync(
        address = ((wpp===0)? (wsCnt @@ wfCnt) | (wfCnt @@ wsCnt)),
        enable  = fire
    )

    /**
     * Last signal propagation for burst boundaries
     *
     * Delays the last signal to match the transposition latency
     * and maintain proper burst synchronization.
     */
    io.q.payload.last := RegNextWhen(start,fire)
}